# 数据库工具类

from sqlalchemy.orm import Session
from .base import SessionLocal

class Database:
    """数据库工具类，提供简单的数据库操作接口"""
    
    def connect(self) -> Session:
        """获取数据库会话"""
        return SessionLocal()
    
    def disconnect(self, session: Session):
        """关闭数据库会话"""
        session.close()
    
    def execute(self, obj, commit=False):
        """执行数据库操作"""
        session = self.connect()
        try:
            session.add(obj)
            if commit:
                session.commit()
                session.refresh(obj)
            return obj
        except Exception as e:
            if commit:
                session.rollback()
            raise e
        finally:
            self.disconnect(session)
    
    def query(self, model_class, *criterion):
        """执行查询操作"""
        session = self.connect()
        try:
            query = session.query(model_class)
            if criterion:
                query = query.filter(*criterion)
            return query.all()
        finally:
            self.disconnect(session)